/*
 * Copyright (c) Huawei Technologies Co., Ltd. 2023-2023. All rights reserved.
 * Description: alg embedding service.
 * Author: qingjicheng
 * Create: 2023-3-2
 */

#include "alg_embedding_service.h"

using namespace std;

namespace hccl {
HcclResult AlgEmbeddingService::AlgServiceCancel(HcclOperator<EmbeddingServiceParam> *opPtr)
{
    std::map<u32, DeviceMem> transferMems;
    if (!opPtr->param_.getRequestParam.transferMems.empty()) {
        for (auto &pair : opPtr->param_.getRequestParam.transferMems) {
            transferMems[pair.first] = pair.second[1];
        }
    }

    CHK_RET(PushSendSignalTask(&cancelSignal_, opPtr->param_.getRequestParam.workerList, transferMems, opPtr,
        OpTaskType::SEND));
    return HCCL_SUCCESS;
}

HcclResult AlgEmbeddingService::PushSendSignalTask(u32 *signal, std::vector<u32> &destRanks,
    std::map<u32, DeviceMem> &transferMems, HcclOperator<EmbeddingServiceParam> *opPtr, OpTaskType sendType)
{
    std::map<u32, void *> rankSendMemMap;
    void *sendMem = signal;
    for (auto &rank : destRanks) {
        if (!transferMems.empty()) {
            OpTaskPtr taskPtr(new (std::nothrow) MemTask());
            MemTask *memTaskPtr = static_cast<MemTask *>(taskPtr.get());
            memTaskPtr->dstMem = transferMems[rank].ptr();
            memTaskPtr->maxSize = transferMems[rank].size();
            memTaskPtr->srcMem = signal;
            memTaskPtr->size = sizeof(u32);
            memTaskPtr->taskType = OpTaskType::MEMCPY;
            rankSendMemMap[rank] = memTaskPtr->dstMem;
            opPtr->PushTask(std::move(taskPtr));
        }
    }

    for (auto &rank : destRanks) {
        TransTask *transTaskPtr = new (std::nothrow) TransTask();
        CHK_PTR_NULL(transTaskPtr);
        transTaskPtr->dstRank = rank;
        transTaskPtr->sendMem = rankSendMemMap[rank] == nullptr ? sendMem : rankSendMemMap[rank];
        transTaskPtr->dataType = HCCL_DATA_TYPE_INT8;
        transTaskPtr->count = sizeof(u32);
        transTaskPtr->outRequest = opPtr->param_.getRequestParam.outRequest;
        transTaskPtr->taskType = sendType;
        transTaskPtr->tag = opPtr->param_.tag;
        transTaskPtr->errorFlag = true;
        transTaskPtr->srcRank = opPtr->param_.userRank;
        OpTaskPtr taskPtr(transTaskPtr);
        opPtr->PushTask(std::move(taskPtr));
    }
    return HCCL_SUCCESS;
}

// 不发送有效数据，只发送空信封
HcclResult AlgEmbeddingService::SendNoSignalTask(vector<u32> &destRanks, map<u32, vector<DeviceMem>> &transferMems,
    HcclOperator<EmbeddingServiceParam> *opPtr, OpTaskType sendType, void *responseValue)
{
    for (auto &rank : destRanks) {
        TransTask *transTaskPtr = new (nothrow) TransTask();
        CHK_PTR_NULL(transTaskPtr);
        OpTaskPtr taskPtr(transTaskPtr);

        transTaskPtr->dstRank = rank;
        transTaskPtr->dataType = HCCL_DATA_TYPE_INT8;
        // worker必须要感知长度，不然误以为未收到ps响应。可优化采用专用控制消息通道。
        transTaskPtr->count = sizeof(u8);
        transTaskPtr->outRequest = opPtr->param_.getRequestParam.outRequest;
        transTaskPtr->taskType = sendType;
        transTaskPtr->tag = opPtr->param_.tag;

        auto iter = transferMems.find(rank);
        if (iter != transferMems.end()) {
            // 需要有效的内存地址，外部统一捕获异常
            transTaskPtr->sendMem = iter->second.at(0).ptr();
        } else {
            // 可能是rdma的UpdateResponse
            transTaskPtr->sendMem = responseValue;
            HCCL_DEBUG("rank[%u] sendMem is responseValue addr[%llu]", rank, hash<void *>{}(responseValue));
        }

        unordered_map<u32, struct RdmaBuffer> &rdmaResponseAddrs = opPtr->param_.getRequestParam.rdmaResponseAddrs;

        auto responseIter = rdmaResponseAddrs.find(rank);
        if (responseIter == rdmaResponseAddrs.end()) {
            HCCL_ERROR("rank[%u] is not found in rdmaResponseAddrs", rank);
            return HCCL_E_INTERNAL;
        }

        transTaskPtr->responseBuffer = responseIter->second;

        HCCL_DEBUG("responseBuffer.addr[%#llx], responseBuffer.buffKey[%#llx]",
            transTaskPtr->responseBuffer.addr, transTaskPtr->responseBuffer.buffKey);

        opPtr->PushTask(move(taskPtr));
    }

    return HCCL_SUCCESS;
}
} // namespace hccl